﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace CoreRay.Samplers
{
    public class StratifiedSampler : Sampler
    {
        private bool jitter;
        private int samplesCountX;
        private int samplesCountY;
        private int samplePosition;
        private int currentPositionX;
        private int currentPositionY;
        private Random random;
        private float[] strataData;

        public StratifiedSampler(RenderRegion renderRegion, int samplesCountX, int samplesCountY, bool jitter) : base(renderRegion, samplesCountX * samplesCountY)
        {
            this.samplesCountX = samplesCountX;
            this.samplesCountY = samplesCountY;
            this.jitter = jitter;

            this.currentPositionX = this.RenderRegion.PositionX;
            this.currentPositionY = this.RenderRegion.PositionY;

            this.random = new Random();

            this.strataData = GenerateStrata(this.samplesCountX, this.samplesCountY, this.jitter);

            for(int i = 0; i < this.samplesCountX * this.samplesCountY * 2; i += 2)
            {
                strataData[i] += this.currentPositionX;
                strataData[i + 1] += this.currentPositionY;
            }   
        }

        protected override Sample GetNextSample()
        {
            if(this.samplePosition == this.samplesCountX * this.samplesCountY)
            {
                this.currentPositionX++;

                if(this.currentPositionX == this.RenderRegion.PositionX + this.RenderRegion.Width)
                {
                    this.currentPositionX = this.RenderRegion.PositionX;
                    this.currentPositionY++;
                }

                if(this.currentPositionY == this.RenderRegion.PositionY + this.RenderRegion.Height)
                {
                    return null;
                }

                strataData = GenerateStrata(this.samplesCountX, this.samplesCountY, this.jitter);

                for(int i = 0; i < this.samplesCountX * this.samplesCountY * 2; i += 2)
                {
                    strataData[i] += this.currentPositionX;
                    strataData[i + 1] += this.currentPositionY;
                }

                this.samplePosition = 0;
            }

            Sample sample = new Sample(strataData[this.samplePosition * 2], strataData[this.samplePosition * 2 + 1]);

            this.samplePosition++;

            return sample;
        }

        private float[] GenerateStrata(int samplesCountX, int samplesCountY, bool jitter)
        {
            float[] strataData = new float[samplesCountX * samplesCountY * 2];

            float dx = 1.0f / samplesCountX;
            float dy = 1.0f / samplesCountY;

            for(int y = 0; y < samplesCountY; y++)
            {
                for(int x = 0; x < samplesCountX; x++)
                {
                    float jitterX = jitter ? RandomFloat() : 0.5f;
                    float jitterY = jitter ? RandomFloat() : 0.5f;
                    
                    strataData[y * samplesCountX * 2 + x * 2] = (x + jitterX) * dx;
                    strataData[y * samplesCountX * 2 + x * 2 + 1] = (y + jitterY) * dy;
                }
            }

            return strataData;
        }

        private float RandomFloat()
        {
            float randomValue = (float) this.random.NextDouble();

            //HACK: Random should be to 1.0f non-inclusive
            if(randomValue > 0.99f)
            {
                randomValue = 0.99f;
            }

            return randomValue;
        }
    }
}
